package com.fliad.ai.integration;

import com.fliad.ai.components.WebInputCom1;
import com.fliad.ai.controller.*;
import com.fliad.ai.entity.AiLlm;
import com.fliad.ai.entity.AiWorkflow;
import org.dromara.autotable.core.AutoTableBootstrap;
import org.noear.dami.Dami;
import org.noear.snack.ONode;
import org.noear.solon.Solon;
import org.noear.solon.core.AppContext;
import org.noear.solon.core.Plugin;
import org.noear.solon.core.event.AppLoadEndEvent;
import org.noear.solon.core.event.EventListener;
import org.noear.solon.core.util.RunUtil;
import org.noear.solon.flow.Chain;
import org.noear.solon.flow.FlowEngine;
import org.noear.solon.web.staticfiles.StaticConfig;
import org.noear.solon.web.staticfiles.StaticMappings;
import org.noear.solon.web.staticfiles.StaticRepository;
import org.noear.solon.web.staticfiles.repository.ClassPathStaticRepository;
import org.noear.wood.DbContext;

import javax.sql.DataSource;
import java.util.List;

/**
 * @author wyl
 * @date 2025年05月04日 下午9:21
 */
public class AiPlugin implements Plugin {

    StaticRepository staticRepository;
    AppContext appContext = null;

    @Override
    public void start(AppContext context) throws Throwable {
        appContext = context;
        if (context.getClassLoader() != Solon.context().getClassLoader()) {
            // 添加自己的静态文件
            // 添加自己的静态文件
            staticRepository = new ClassPathStaticRepository(context.getClassLoader(), StaticConfig.RES_STATIC_LOCATION);
            StaticMappings.add("/", staticRepository);
            Thread.currentThread().setContextClassLoader(context.getClassLoader());
            AutoTableBootstrap.start();
            Thread.currentThread().setContextClassLoader(Solon.context().getClassLoader());
        }
        context.onEvent(AppLoadEndEvent.class, new EventListener<AppLoadEndEvent>() {
            @Override
            public void onEvent(AppLoadEndEvent appLoadEndEvent) throws Throwable {
                Dami.bus().send("add_no_login_path", "api/chat/**");
                Dami.bus().send("add_no_login_path", "ai/llm");
                Dami.bus().send("add_no_login_path", "api/task/**");
                Dami.bus().send("add_no_login_path", "ai/workflow/**");
            }
        });
        context.beanScan("com.fliad.ai.service");
        context.beanScan("com.fliad.ai.listener");
        context.beanScan("com.fliad.ai.components");
        context.beanScan(WebInputCom1.class);
        context.app().add("/ai/workflow", AiWorkflowController.class);
        context.app().add("/ai/llm", AiLlmController.class);
        context.app().add("/ai/chat", AiChatController.class);
        // 新增知识库控制器路由
        context.app().add("/ai/knowledge", AiKnowledgeController.class);
        // 新增文档控制器路由
        context.app().add("/ai/document", AiDocumentController.class);
        // 新增MCP控制器路由
        context.app().add("/ai/mcp", AiMcpController.class);
        // 新增TOOL控制器路由
        context.app().add("/ai/tool", ToolController.class);
        context.app().add("/api/task", TaskController.class);
        // 新增智能体控制器路由
        context.app().add("/ai/agent", AiAgentController.class);

        context.subBeansOfType(DataSource.class, dataSource -> {
            DbContext dbContext = new DbContext(dataSource);
            try {
                List<AiWorkflow> list = dbContext.table("ai_workflow").selectList("*", AiWorkflow.class);

                for (int i = 0; i < list.size(); i++) {
                    ONode node = ONode.newObject();
                    node.set("id", list.get(i).getId());
                    int finalI = i;
                    node.getOrNew("layout").asArray().build(layoutNode -> {
                        ONode content = ONode.loadStr(list.get(finalI).getContent());
                        content.get("nodes").forEach(taskNode -> {
                            ONode node_ = ONode.newObject().set("id", taskNode.get("id"));
                            // 开始节点
                            if ("startNode".equals(taskNode.get("type").getString())) {
                                node_.set("task", "@WebInput");
                            } else if ("llmNode".equals(taskNode.get("type").getString())) {
                                node_.set("task", "@ChatModel");
                                ONode meta = node_.getOrNew("meta");
                                meta.set("systemPrompt", taskNode.get("init_data").get("systemPrompt"));
                                try {
                                    AiLlm aiLlm = dbContext.table("ai_llm").whereEq("id", taskNode.get("init_data").get("llmId").getString())
                                            .selectItem("*", AiLlm.class);
                                    meta.getOrNew("chatConfig")
                                            .set("provider", aiLlm.getDialect())
                                            .set("model", aiLlm.getLlmModel())
                                            .set("apiUrl", aiLlm.getLlmEndpoint())
                                            .set("apiKey", aiLlm.getLlmApiKey());
                                } catch (Exception e) {
                                    // ...
                                }

                            } else if ("endNode".equals(taskNode.get("type").getString())) {
                                node_.set("task", "@WebOutput");
                            }
                            node_.set("id", taskNode.get("id"));
                            layoutNode.addNode(node_);
                        });
                    });

                    RunUtil.delay(() -> {
                        FlowEngine flowEngine = Solon.context().getBean(FlowEngine.class);

                        flowEngine.load(Chain.parseByDom(node));
                    }, 1000);
                }
            } catch (Exception e) {
                // ...
            }
        });
    }

    @Override
    public void stop() throws Throwable {
        //移除静态文件仓库
        if (staticRepository != null) {
            StaticMappings.remove(staticRepository);
            staticRepository = null;
        }
        appContext.app().router().remove("/ai/workflow");
        appContext.app().router().remove("/ai/llm");
        appContext.app().router().remove("/ai/chat");
        // 新增移除知识库控制器路由
        appContext.app().router().remove("/ai/knowledge");
        // 新增移除文档控制器路由
        appContext.app().router().remove("/ai/document");
        // 新增移除MCP控制器路由
        appContext.app().router().remove("/ai/mcp");
        appContext.clear();
    }

}
